package org.example.datasource;

import com.alibaba.druid.pool.DruidDataSource;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.RemovalCause;
import com.github.benmanes.caffeine.cache.RemovalListener;
import lombok.extern.slf4j.Slf4j;
import org.example.utils.CommonUtils;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantReadWriteLock;

@Slf4j
public class DynamicDataSource extends AbstractRoutingDataSource {
    // 读写锁用于保护共享状态
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock();
    private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock();

    // Caffeine缓存，最大100个数据源，访问后30分钟过期
    private final Cache<String, DataSource> dataSourceCache = Caffeine.newBuilder()
            .maximumSize(3)
            .expireAfterAccess(10, TimeUnit.SECONDS)
            .removalListener((RemovalListener<String, DataSource>) (key, value, cause) -> {
                if (cause == RemovalCause.SIZE) {
                    log.info("Removing data source from cache: {}, cause: {}", key, cause);
                    removeDataSource(key);
                }
            })
            .build();

    @Override
    protected Object determineCurrentLookupKey() {
        readLock.lock();
        try {
            String dsKey = SubjectContext.getDataSourceKey();
            if (dsKey == null) {
                return CommonUtils.getPrimaryDsKey();
            }
            //触发使用次数 + 时间
            DataSource dataSource = dataSourceCache.getIfPresent(dsKey);
            if (dataSource != null) {
                return dsKey;
            }
            return dsKey;
        } finally {
            readLock.unlock();
        }
    }

    // 动态添加数据源
    public void addDataSource(String dsKey, DataSource dataSource) {
        writeLock.lock();
        try {
            log.debug("Adding new data source: {}", dsKey);
            Map<Object, Object> targetDataSources = new HashMap<>(getResolvedDataSources());
            targetDataSources.put(dsKey, dataSource);

            setTargetDataSources(targetDataSources);

            afterPropertiesSet(); // 刷新数据源

            // 如果是主数据源，不放入缓存
            if (!CommonUtils.getPrimaryDsKey().equals(dsKey)) {
                dataSourceCache.put(dsKey, dataSource);
            }
        } finally {
            writeLock.unlock();
        }
    }

    // 删除数据源
    public void removeDataSource(String dsKey) {
        writeLock.lock();
        try {
            log.debug("Removing data source: {}", dsKey);
            //1. 从缓存移除（触发 RemovalListener）
            if (dataSourceCache.getIfPresent(dsKey) != null) {
                dataSourceCache.invalidate(dsKey);
            }

            Map<Object, Object> targetDataSources = new HashMap<>(getResolvedDataSources());
            Object object = targetDataSources.get(dsKey);
            if (object == null) {
                return;
            }
            //移除数据源
            targetDataSources.remove(dsKey);
            //关闭数据源
            DataSource oldDataSource = (DataSource) object;
            // 安全关闭旧数据源
            if (oldDataSource instanceof DruidDataSource) {
                try {
                    ((DruidDataSource) oldDataSource).close();
                    log.info("Successfully closed data source: {}", dsKey);
                } catch (Exception e) {
                    log.error("Failed to close data source: {}", dsKey, e);
                }
            }

            setTargetDataSources(targetDataSources);

            //刷新数据源
            afterPropertiesSet();
        } finally {
            writeLock.unlock();
        }
    }
}